﻿#include <QDebug>

#include "AVFileStoreMP3.h"

static int check_sample_fmt(const AVCodec *codec, enum AVSampleFormat sample_fmt)
{
    const enum AVSampleFormat *p = codec->sample_fmts;

    while (*p != AV_SAMPLE_FMT_NONE) {
        if (*p == sample_fmt)
            return 1;
        p++;
    }
    return 0;
}

AVFileStoreMP3::AVFileStoreMP3(const QString &fileName) :AVFileStoreAbstract(), m_fileName(fileName)
{
    m_defaultSampleRate = 44100;
    m_defaultChannels = 2;
    m_defaultFormat = (int)AV_SAMPLE_FMT_S16P;
}

bool AVFileStoreMP3::init()
{
    int bitRate = 128000;
    int channelLayout = (int)AV_CH_LAYOUT_STEREO;

    AVSampleFormat format = AV_SAMPLE_FMT_S16P;
    //编码器
    const AVCodec *codecMp3;
    //编码器上下文
    AVCodecContext *contexMp3= NULL;
    //编码后的mp3数据包
    /* 找到MP3的编码器 */
    codecMp3 = avcodec_find_encoder(AV_CODEC_ID_MP3);
    if (!codecMp3) {
        fprintf(stderr, "Codec not found\n");
        exit(1);
        return false;
    }

    //生成编码器上下文
    contexMp3 = avcodec_alloc_context3(codecMp3);
    if (!contexMp3) {
        fprintf(stderr, "Could not allocate audio codec context\n");
        exit(1);
        return false;
    }
    /* 设置编码器上下文的bit率 */
    contexMp3->bit_rate = bitRate;
    contexMp3->sample_fmt = format;
    if (!check_sample_fmt(codecMp3, contexMp3->sample_fmt)) {
        fprintf(stderr, "Encoder does not support sample format %s",
                av_get_sample_fmt_name(contexMp3->sample_fmt));
        exit(1);
        return false;
    }
    /* 设置编码器的必要参数 */
    contexMp3->sample_rate    = m_defaultSampleRate;
    contexMp3->channel_layout = channelLayout;
    contexMp3->channels       = av_get_channel_layout_nb_channels(contexMp3->channel_layout);
    if (avcodec_open2(contexMp3, codecMp3, NULL) < 0) {
        fprintf(stderr, "Could not open codec\n");
        exit(1);
        return false;
    }

    m_codecCtx = contexMp3;

    m_mp3File = fopen(m_fileName.toStdString().c_str(),"wb");
    if (!m_mp3File) {
        return false;
    }

    return true;
}

bool AVFileStoreMP3::destory()
{
    fclose(m_mp3File);
    avcodec_free_context(&m_codecCtx);
    return true;
}

bool AVFileStoreMP3::fillFrame(AVFrame *frame)
{
    static AVPacket *pktMp3_Out = av_packet_alloc(); /* 释放问题 */

    uint8_t buf[4096] = {0};// 1024 * 2 * 2
    int retBytes = 0;
    int retNBsample = 0;

    AVFrame * resampleFrame = 0;

    if (frame->sample_rate != m_defaultSampleRate ||
            frame->channels != m_defaultChannels ||
            frame->format != m_defaultFormat) {
        if (resample(frame,buf,&retBytes,&retNBsample) ) {
            resampleFrame = av_frame_alloc();
            resampleFrame->sample_rate = m_defaultSampleRate;
            resampleFrame->channels = m_defaultChannels;
            resampleFrame->format = m_defaultFormat;
            resampleFrame->nb_samples = retNBsample;
            if (avcodec_fill_audio_frame(resampleFrame,m_defaultChannels,(AVSampleFormat)m_defaultFormat,buf,retBytes,1) < 0) {
                qDebug() << "fill error";
            }

        } else {
            qDebug() << "resample error";
        }
    } else {
        resampleFrame = av_frame_clone(frame);
    }

    /* 编码frame */
    int ret = avcodec_send_frame(m_codecCtx, resampleFrame);
    if (ret < 0) {
        fprintf(stderr, "Error sending the frame to the encoder\n");
        exit(1);
    }

    av_frame_unref(resampleFrame);
    av_frame_free(&resampleFrame);

    /* 读取编码出的数据包，通常不止一个 */
    while (ret >= 0) {
        ret = avcodec_receive_packet(m_codecCtx, pktMp3_Out);
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
            return true;
        } else if (ret < 0) {
            fprintf(stderr, "Error encoding audio frame\n");
            return false;
        } else {
            static long flushSize = 0;
            fwrite(pktMp3_Out->data, 1, pktMp3_Out->size, m_mp3File);
            flushSize >= 5120 ? fflush(m_mp3File) : flushSize += pktMp3_Out->size;
            av_packet_unref(pktMp3_Out);
        }
    }
    return true;
}
